'''
Created on Mar 1, 2020
Pytorch Implementation of LightGCN in
Xiangnan He et al. LightGCN: Simplifying and Powering Graph Convolution Network for Recommendation

@author: Jianbai Ye (gusye@mail.ustc.edu.cn)
'''
import world
import torch
from torch import nn, optim
import numpy as np
from torch import log
from dataloader import BasicDataset
from time import time
from model import LightGCN
from model import PairWiseModel
from sklearn.metrics import roc_auc_score
import random
import os
try:
    from cppimport import imp_from_filepath
    from os.path import join, dirname
    path = join(dirname(__file__), "sources/sampling.cpp")
    sampling = imp_from_filepath(path)
    sampling.seed(world.seed)
    sample_ext = True
except:
    world.cprint("Cpp extension not loaded")
    sample_ext = False

#计算bpr损失函数
class BPRLoss:
    #recmodel指定类型为 PairWiseModel config指定类型为字典
    def __init__(self,
                 recmodel : PairWiseModel,
                 config : dict):
        self.model = recmodel
        #decay l2正则化惩罚力度 default=1e-4
        self.weight_decay = config['decay']
        self.lr = config['lr']
        self.opt = optim.Adam(recmodel.parameters(), lr=self.lr)

    #stageOne即开始训练反向传播 users, pos, neg都是一维tensor [2048]
    def stageOne(self, users, pos, neg):
        #损失函数loss,reg_loss分别为bpr损失和正则化损失
        loss, reg_loss = self.model.bpr_loss(users, pos, neg)
        reg_loss = reg_loss*self.weight_decay
        loss = loss + reg_loss

        #反向传播
        self.opt.zero_grad()
        loss.backward()
        self.opt.step()
        #将计算出的损失从 GPU 转移到 CPU，并返回损失值。这是为了方便后续处理或日志记录。
        return loss.cpu().item()

#返回一个numpy类型 [users,pos,neg] 每个用户随机选交互过的一个商品和没有交互过的一个商品
def UniformSample_original(dataset, neg_ratio = 1):
    dataset : BasicDataset
    #返回用户交互过的商品 是一个列表 例allPos[i,:]获取第i行数据 即用户i所有交互过的商品列表
    allPos = dataset.allPos
    start = time()
    #sample_ext:c++抽样函数是否可用
    if sample_ext:
        #c++实现
        S = sampling.sample_negative(dataset.n_users, dataset.m_items,
                                     dataset.trainDataSize, allPos, neg_ratio)
    else:
        #python实现
        S = UniformSample_original_python(dataset)
    return S
#python实现:返回一个numpy类型 [users,pos,neg] 每个用户随机选交互过的一个商品和没有交互过的一个商品
def UniformSample_original_python(dataset):
    """
    the original impliment of BPR Sampling in LightGCN
    :return:
        np.array
    """
    total_start = time()
    dataset : BasicDataset
    user_num = dataset.trainDataSize
    #生成user_num个随机数 每个随机数范围为[0,dataset.n_users-1]
    users = np.random.randint(0, dataset.n_users, user_num)
    #dataset.allPos返回用户交互过的商品 是一个列表 例allPos[i,:]获取第i行数据 即用户i所有交互过的商品列表
    allPos = dataset.allPos
    S = []
    sample_time1 = 0.
    sample_time2 = 0.
    #所有用户 每个用户选一个交互过的商品和没交互过的商品
    for i, user in enumerate(users):
        start = time()
        #获取用户交互过的商品
        posForUser = allPos[user]
        if len(posForUser) == 0:
            continue
        sample_time2 += time() - start
        #随机选一个用户交互过的商品
        posindex = np.random.randint(0, len(posForUser))
        positem = posForUser[posindex]
        #随机选一个用户没有交互过的商品
        while True:
            negitem = np.random.randint(0, dataset.m_items)
            if negitem in posForUser:
                continue
            else:
                break
        #一个用户交互过的商品和没交互过的商品
        S.append([user, positem, negitem])
        end = time()
        sample_time1 += end - start
    total = time() - total_start
    return np.array(S)

# ===================end samplers==========================
# =====================utils====================================
#设置随机数种子
def set_seed(seed):
    np.random.seed(seed)
    if torch.cuda.is_available():
        #设置了 CUDA 随机数种子后，可以确保在使用 CUDA 加速时生成的随机数是可复现的，即在相同种子下，随机数生成的结果是确定的。
        torch.cuda.manual_seed(seed)
        torch.cuda.manual_seed_all(seed)
    torch.manual_seed(seed)

#选择模型路径
def getFileName():
    if world.model_name == 'mf':
        file = f"mf-{world.dataset}-{world.config['latent_dim_rec']}.pth.tar"
    elif world.model_name == 'lgn':
        file = f"lgn-{world.dataset}-{world.config['lightGCN_n_layers']}-{world.config['latent_dim_rec']}.pth.tar"
    return os.path.join(world.FILE_PATH,file)
#*把所有参数存放到元组 **把参数存到字典
"""
minibatch这个函数允许用户传入一个或多个张量，并返回这些张量的小批量数据，每次返回一个小批量数据。
这对于在机器学习或深度学习中处理大量数据非常有用，因为它可以减少内存使用并提高处理速度。
"""
def minibatch(*tensors, **kwargs):
    #从关键字参数中获取 batch_size 的值，如果未提供，则默认为 world.config['bpr_batch_size']
    batch_size = kwargs.get('batch_size', world.config['bpr_batch_size'])
    if len(tensors) == 1:
        tensor = tensors[0]
        for i in range(0, len(tensor), batch_size):
            yield tensor[i:i + batch_size]
    else:
        #我们主要输入users posItems negItems 这三个总长度是一样的所以取第一个即可len(tensors[0])
        for i in range(0, len(tensors[0]), batch_size):
            #对于每个小批量，从每个张量中提取相应的小批量，并返回一个元组。同样，这也是使用生成器的方式返回数据。
            #x遍历整个tensors数组 x为依次users posItems negItems 对上述三个取minibatch
            yield tuple(x[i:i + batch_size] for x in tensors)

#输入都为[user_len,1]
def shuffle(*arrays, **kwargs):

    require_indices = kwargs.get('indices', False)
    #len(x) for x in arrays或许数组所有元素的长度
    #set(len(x) for x in arrays)将上一步结果放入集合即去重
    #len(set(len(x) for x in arrays))查询集合长度如果不为1则说明数组中有的元素长度不一致
    if len(set(len(x) for x in arrays)) != 1:
        raise ValueError('All inputs to shuffle must have '
                         'the same length.')
    #输入的arrays中array都为[user_len,1] 所以取第一个就可以了即array[0] 获取其长度即np.arange(user_len)生成一个长为user_len的数组
    shuffle_indices = np.arange(len(arrays[0]))
    #np.random.shuffle 是 NumPy 库中的一个函数，用于随机打乱一个数组的顺序。
    np.random.shuffle(shuffle_indices)

    if len(arrays) == 1:
        result = arrays[0][shuffle_indices]
    else:
        #长度大于1说明arrays中所有元素都要随机排列
        #依次访问shuffle_indices中的数字
        #例shuffle_indices=[0,1,2] 则x[shuffle_indices] for x in arrays会生成一个数组[x[0],x[1],x[2]]
        result = tuple(x[shuffle_indices] for x in arrays)
    #在python中0解释为false 非0解释为true
    if require_indices:
        return result, shuffle_indices
    else:
        return result


class timer:
    """
    Time context manager for code block
        with timer():
            do something
        timer.get()
    """
    from time import time
    TAPE = [-1]  # global time record
    #NAMED_TAPE字典
    NAMED_TAPE = {}

    @staticmethod
    #获取最新的时间
    def get():
        if len(timer.TAPE) > 1:
            return timer.TAPE.pop()
        else:
            return -1

    @staticmethod
    #参数select_keys被设置为None，这意味着如果调用者没有提供该参数，它默认会是None
    def dict(select_keys=None):
        hint = "|"
        if select_keys is None:
            for key, value in timer.NAMED_TAPE.items():
                hint = hint + f"{key}:{value:.2f}|"
        else:
            for key in select_keys:
                value = timer.NAMED_TAPE[key]
                hint = hint + f"{key}:{value:.2f}|"
        return hint

    @staticmethod
    #全部置为0
    def zero(select_keys=None):
        if select_keys is None:
            for key, value in timer.NAMED_TAPE.items():
                timer.NAMED_TAPE[key] = 0
        else:
            for key in select_keys:
                timer.NAMED_TAPE[key] = 0
    #**kwargs接受任意参数 最后结果为一个字典
    def __init__(self, tape=None, **kwargs):
        if kwargs.get('name'):
            #给timer.NAMED_TAPE[kwargs['name']]一个初始值 如果存在则直接赋值 不存在给0
            #例输入timer(name='Sample') 则下面语句最终结果为timer.NAMED_TAPE[Sample] = 0
            timer.NAMED_TAPE[kwargs['name']] = timer.NAMED_TAPE[
                kwargs['name']] if timer.NAMED_TAPE.get(kwargs['name']) else 0.
            self.named = kwargs['name']
            if kwargs.get("group"):
                #TODO: add group function
                pass
        else:
            self.named = False
            self.tape = tape or timer.TAPE

    #记录开始时间
    def __enter__(self):
        self.start = timer.time()
        return self
    #记录结束时间
    def __exit__(self, exc_type, exc_val, exc_tb):
        if self.named:
            timer.NAMED_TAPE[self.named] += timer.time() - self.start
        else:
            self.tape.append(timer.time() - self.start)


# ====================Metrics==============================
# =========================================================
#获取recall和precision
"""
test_data是测试集用户交互过的商品列表 其形状为[batch][测试集实际交互的商品数量] 这里batch为100
r为实际数据和预测数据的命中情况每个元素命中为1.没命中为0. 形状为[batch][k] batch为一个批次的大小 k为排前k个商品 在这里默认为[100][20] 
k 默认为20
"""
#计算recall和precis 返回一个字典 字典中每个value值都为一个标量 {string:float}
def RecallPrecision_ATk(test_data, r, k):
    #选择所有行和0~k-1列 sum(1)即二维数组沿着第一维求和 即[100,20]->[100]
    #统计每个用户的命中次数 right_pred形状为[batch] 即[100]
    right_pred = r[:, :k].sum(1)
    #每个用户预测k个商品
    precis_n = k
    #len(test_data)获取一组用户 len(test_data[i])获取一个用户交互过的所有商品数量 recall_n形状为[batch] 即[100]
    recall_n = np.array([len(test_data[i]) for i in range(len(test_data))])
    #对应元素相除 再把所有结果相加 因为是一维数组所以最终结果recall为一个标量
    recall = np.sum(right_pred/recall_n)
    #对应元素相除 再把所有结果相加 因为是一维数组所以最终结果precis为一个标量
    precis = np.sum(right_pred)/precis_n
    #返回一个字典其中每个value值都为一个标量
    return {'recall': recall, 'precision': precis}

"""
r为实际数据和预测数据的命中情况每个元素命中为1.没命中为0. 形状为[batch][k] batch为一个批次的大小 k为排前k个商品 在这里默认为[100][20] 
k 默认为20
"""
#MRR分数越高越好 输出一个标量
def MRRatK_r(r, k):
    # 统计每个用户的命中次数 right_pred形状为[batch] 即[100]
    pred_data = r[:, :k]
    # 排名越靠前log(排名)越小 分母越小值越大 即排名靠前的命中权值高
    scores = np.log2(1./np.arange(1, k+1))
    #将所有命中数据和没命中数据除以scores
    pred_data = pred_data/scores
    #对每个用户的数据求和
    pred_data = pred_data.sum(1)
    return np.sum(pred_data)
"""
test_data是测试集用户交互过的商品列表 其形状为[batch][测试集实际交互的商品数量] 这里batch为100
r为实际数据和预测数据的命中情况每个元素命中为1.没命中为0. 形状为[batch][k] batch为一个批次的大小 k为排前k个商品 在这里默认为[100][20] 
k 默认为20
输出一个batch中所有用户的nbcg得分 是一个标量
"""
def NDCGatK_r(test_data,r,k):
    """
    Normalized Discounted Cumulative Gain
    rel_i = 1 or 0, so 2^{rel_i} - 1 = 1 or 0
    """
    assert len(r) == len(test_data)
    # 统计每个用户的命中次数 right_pred形状为[batch] 即[100]
    pred_data = r[:, :k]
    #创建一个全为0的矩阵 其维度为[batch][k] 默认为[100][20]
    test_matrix = np.zeros((len(pred_data), k))
    for i, items in enumerate(test_data):
        #选k和用户的实际交互商品数量中较短的那个
        length = k if k <= len(items) else len(items)
        #把第i行0~length-1列都设置为1
        test_matrix[i, :length] = 1
    #max_r其维度为[batch][k] 默认为[100][20] 其中有数据的部分为1 每行数据不足k的地方为0
    max_r = test_matrix
    """ 
    idcg为理想分数即每个商品都成功预测的分数所以用max_r矩阵
    np.arange(2, k + 2): 创建一个从2到k+1的数组。ndcg中第一个元素是除以1 所以从第二个元素开始即可 第一个元素保持原状态
    max_r * 1./np.log2(np.arange(2, k + 2)) [batch][k]除以[k]使用广播机制 每行都和k维度数组相除
    axis=1 按行相加即[batch][k] sum后 变为 [batch] 计算每个用户的得分
    """
    idcg = np.sum(max_r * 1./np.log2(np.arange(2, k + 2)), axis=1)
    """
    dcg为实际分数 pred_data为真实的预测成功次数
    np.arange(2, k + 2): 创建一个从2到k+1的数组。ndcg中第一个元素是除以1 所以从第二个元素开始即可 第一个元素保持原状态
    max_r * 1./np.log2(np.arange(2, k + 2)) [batch][k]除以[k]使用广播机制 每行都和k维度数组相除
    axis=1 按行相加即[batch][k] sum后 变为 [batch] 计算每个用户的得分
    """
    dcg = pred_data*(1./np.log2(np.arange(2, k + 2)))
    dcg = np.sum(dcg, axis=1)
    #在理想分数中为0的分数设置为1 因为idcg要作为被除数要求不能为0
    idcg[idcg == 0.] = 1.
    ndcg = dcg/idcg
    #把原本为NaN的值替换为0
    ndcg[np.isnan(ndcg)] = 0.
    #输出所有用户相加的nbcg值
    return np.sum(ndcg)

"""
all_item_scores为单个用户的预测评分[1][item_len]
dataset为输入的数据集
test_data
"""

def AUC(all_item_scores, dataset, test_data):
    """
        design for a single user
    """
    dataset : BasicDataset
    r_all = np.zeros((dataset.m_items, ))
    r_all[test_data] = 1
    r = r_all[all_item_scores >= 0]
    test_item_scores = all_item_scores[all_item_scores >= 0]
    return roc_auc_score(r, test_item_scores)

#获取预测数据和测试数据矩阵 返回[batch][k] batch为一个批次的大小 k为排前k个商品 在这里默认为[100][20] 每个元素命中为1.没命中为0.
def getLabel(test_data, pred_data):
    r = []
    #这里的len(test_data)长为100 默认是每个batch为100
    for i in range(len(test_data)):
        groundTrue = test_data[i]
        predictTopK = pred_data[i]
        #将 predictTopK 中的每个元素与 groundTrue 进行比较，检查它是否在 groundTrue 中。最后返回一个列表里面元素为bool类型
        pred = list(map(lambda x: x in groundTrue, predictTopK))
        #先转为numpy的array类型再把bool转为float 把True转换为浮点数1.0。如果你使用False作为输入，结果将是0.0。
        pred = np.array(pred).astype("float")
        r.append(pred)
    #把r列表转为numpy的array类型再转为float类型
    return np.array(r).astype('float')

# ====================end Metrics=============================
# =========================================================
